import torch
import torch.nn.functional as F
import numpy as np

class SemiLoss(object):
    # mixmatch 文献中 lambda_u 使用的为100
    def __init__(self, rampup_epoches, lambda_u=75) -> None:
        self.lambda_u = lambda_u
        self.rampup_epoches = rampup_epoches
        
    def __call__(self, outputs_x, targets_x, outputs_u, targets_u, epoch):
        probs_u = torch.softmax(outputs_u, dim=1)

        Lx = -torch.mean(torch.sum(F.log_softmax(outputs_x, dim=1) * targets_x, dim=1))
        Lu = torch.mean((probs_u - targets_u)**2)

        return Lx, Lu, self.lambda_u * linear_rampup(epoch, self.rampup_epoches)

def linear_rampup(current, rampup_length=50):
    if rampup_length == 0:
        return 1.0
    else:
        current = np.clip(current / rampup_length, 0.0, 1.0)
        return float(current)